
import numpy
import scipy.linalg
from functools import reduce
from scipy.io import FortranFile
from pyscf import lib, scf, symm

def canonical_rhf(mol, mo_coeff, mo_occ, savefvo = True):

    myhf = scf.RHF(mol)
    dm0 = scf.hf.make_rdm1(mo_coeff, mo_occ)
    fock0 = myhf.get_hcore() + myhf.get_veff(mol, dm0)
    eref = myhf.energy_elec(dm0)[0] + myhf.energy_nuc()

    nmo = len(mo_occ)
    if not mol.symmetry:
       nirrep = 1
    else :
       nirrep = len(mol.irrep_id)
    if nirrep == 1:
       symid = numpy.zeros(nmo, dtype = int)
       mo_coeff = lib.tag_array(mo_coeff, orbsym=symid)
    
    newmo_e, newmo = scf.hf_symm.canonicalize(myhf, mo_coeff, mo_occ, fock=fock0)

    if nirrep>1 :
       symid = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, newmo)
       symid = symid % 10

    occidx = mo_occ == 2
    viridx = mo_occ == 0
    listpop = []
    listvrt = []
    for irrep in range(nirrep):
       idx0 = symid == irrep
       idx1 = occidx & idx0; idx2 = viridx & idx0
       listpop += list(numpy.arange(nmo)[idx1])
       listvrt += list(numpy.arange(nmo)[idx2])

    focknew = newmo.T.dot(fock0).dot(newmo)

    nocc = len(listpop); nvir = len(listvrt)
    print('nocc, nvir', nocc, nvir, nmo)

    fockoo = focknew[listpop][:, listpop]
    fockvv = focknew[listvrt][:, listvrt]
    
    if numpy.amax(fockoo - numpy.diag(numpy.diag(fockoo))) > 1.e-8:
       print('error in diagonal nature of fockoo')
    
    if numpy.amax(fockvv - numpy.diag(numpy.diag(fockvv))) > 1.e-8:
       print('error in diagonal nature of fockvv')

    if savefvo:
       fockvo = focknew[listvrt][:, listpop]
       f = FortranFile('FAIAAR.dat', 'w')
       f.write_record(fockvo.T)

    myhf.mo_occ = mo_occ
    myhf.mo_coeff = newmo
    myhf.mo_energy = newmo_e
    myhf.e_tot = eref

    return myhf

def canonical_rohf(myhf, mol, mo_coeff, mo_occ, nf, savefvo = True):
    dm0 = scf.rohf.make_rdm1(mo_coeff, mo_occ)
    fock0 = myhf.get_hcore() + myhf.get_veff(mol, dm0)
    eref = myhf.energy_elec(dm0)[0] + myhf.energy_nuc()
    print('dm0', dm0.shape)
    print('fock0', fock0.shape)
    print('eref', eref)
#========alpha=========================
    occidxa = mo_occ == 2
    occidxa += mo_occ == 1
    viridxa = mo_occ == 0
    nocca = numpy.count_nonzero(occidxa)
    nvira = numpy.count_nonzero(viridxa)
    print('len of occidx', numpy.count_nonzero(occidxa))
    print('len of viridx', numpy.count_nonzero(viridxa))
    newmoa = numpy.empty_like(mo_coeff)
    newmo_ea = numpy.empty(mo_occ.size)
    for idx in (occidxa, viridxa):
#    print('idx', idx)
        if numpy.count_nonzero(idx) > 0:
#           print('idx2', numpy.count_nonzero(idx))
           orb = mo_coeff[:,idx]
#           print('orb.shape', orb.shape)
           f1 = reduce(numpy.dot, (orb.conj().T, fock0[0,:,:], orb))
           e, c = scipy.linalg.eigh(f1)
           newmoa[:,idx] = numpy.dot(orb, c)
           newmo_ea[idx] = e
    print('ea', newmo_ea)

    focknewa = newmoa.T.dot(fock0[0,:,:]).dot(newmoa)
    
    nmo = len(mo_occ)
#   if not mol.symmetry:
#      nirrep = 1
#   else :
#      nirrep = len(mol.irrep_id)
#   if nirrep == 1:
#      symid = numpy.zeros(nmo, dtype = int)
#   if nirrep>1 :
#      symid = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, newmoa)
#      symid = symid % 10
#   listpopa = []
#   listvrta = []
#   for irrep in range(nirrep):
#      idx0 = symid == irrep
#      idx1 = occidxa & idx0; idx2 = viridxa & idx0
#      listpopa += list(numpy.arange(nmo)[idx1])
#      listvrta += list(numpy.arange(nmo)[idx2])


#   nocca = len(listpopa); nvira = len(listvrta)
#   print('nocca, nvira', nocca, nvira, nmo)

    fockooaa = focknewa[:nocca, :nocca]
    fockvvaa = focknewa[nocca:nmo, nocca:nmo]
#   fockooaa = focknewa[listpopa][:, listpopa]
#   fockvvaa = focknewa[listvrta][:, listvrta]
    
    if numpy.amax(fockooaa - numpy.diag(numpy.diag(fockooaa))) > 1.e-8:
       print('error in diagonal nature of fockooaa')
    
    if numpy.amax(fockvvaa - numpy.diag(numpy.diag(fockvvaa))) > 1.e-8:
       print('error in diagonal nature of fockvvaa')

    if savefvo:
#      fockvoaa = focknewa[listvrta][:, listpopa]
#      print('fockvoaa', fockvoaa.shape) 
       fockvoaa = focknewa[nocca:nmo, nf:nocca]
#      print('fockvoaa', fockvoaa.shape) 
       for i in range(10):
            print('fai',i, fockvoaa[i,0])
      
       f = FortranFile('FAIAAR.dat', 'w')
       f.write_record(fockvoaa.T)
#       print(numpy.linalg.norm(fockvoaa))
#========beta=========================
    occidxb = mo_occ == 2
    viridxb = mo_occ == 1
    viridxb += mo_occ == 0
    noccb = numpy.count_nonzero(occidxb)
    nvirb = numpy.count_nonzero(viridxb)
    newmob = numpy.empty_like(mo_coeff)
    newmo_eb = numpy.empty(mo_occ.size)
    for idx in (occidxb, viridxb):
        if numpy.count_nonzero(idx) > 0:
           orb = mo_coeff[:,idx]
           f1 = reduce(numpy.dot, (orb.conj().T, fock0[1,:,:], orb))
           e, c = scipy.linalg.eigh(f1)
           newmob[:,idx] = numpy.dot(orb, c)
           newmo_eb[idx] = e
    print('eb', newmo_eb)

    focknewb = newmob.T.dot(fock0[1,:,:]).dot(newmob)
    
    nmo = len(mo_occ)

    print('noccb, nvirb', noccb, nvirb, nmo)

    fockoobb = focknewb[:noccb, :noccb]
    fockvvbb = focknewb[noccb:nmo, noccb:nmo]
#   fockooaa = focknewa[listpopa][:, listpopa]
#   fockvvaa = focknewa[listvrta][:, listvrta]
    
    if numpy.amax(fockoobb - numpy.diag(numpy.diag(fockoobb))) > 1.e-8:
       print('error in diagonal nature of fockooaa')
    
    if numpy.amax(fockvvbb - numpy.diag(numpy.diag(fockvvbb))) > 1.e-8:
       print('error in diagonal nature of fockvvaa')

    if savefvo:
#       fockvoaa = focknewa[listvrta][:, listpopa]
       fockvobb = focknewb[noccb:nmo, nf:noccb]
       f = FortranFile('FAIBBR.dat', 'w')
       f.write_record(fockvobb.T)
#       print(numpy.linalg.norm(fockvobb))
    mo_coeff_new = numpy.stack((newmoa,newmob), axis = 0)
    mo_energy_new = numpy.stack((newmo_ea,newmo_eb), axis = 0)
    print('mo_coeff_new', mo_coeff_new.shape)
    print('mo_energy_new', mo_energy_new.shape)
#   print('mo_energy_newa', mo_energy_new[0,:])
#   print('mo_energy_newb', mo_energy_new[1,:])
#   print('mo_energy', myhf.mo_energy)
    myhf.mo_coeff = lib.tag_array(myhf.mo_coeff, ro=mo_coeff_new)
    myhf.mo_energy = lib.tag_array(myhf.mo_energy, ro=mo_energy_new)
#   print('mo_energy', myhf.mo_energy)
    return myhf, mo_coeff_new, mo_energy_new
